import os
import json

from django.db.models import Model
from django.conf import settings
from whoosh.index import EmptyIndexError
from whoosh.analysis import RegexAnalyzer
from whoosh.filedb.filestore import FileStorage
from whoosh.highlight import HtmlFormatter, WholeFragmenter
from whoosh.fields import Schema, TEXT, KEYWORD, ID


# 这个地方灵活性太强，建议参考流程，重新写
class WhooshSearch:
    # 对应的具体的模型
    model: Model = None
    # 需要存储到索引中的键
    keys: list = None
    # 用于查询的键，如果不指定和keys一致（去掉id）
    index_keys: list = None
    # 索引名称
    index_name: str = None

    def __init__(self, query=None, queryset=None):
        """
        创建索引，动态插入索引
        :param query: 查询范围，获取此次需要插入的值
        :param queryset: 直接指定需要插入的值
        """
        if self.model is None or self.keys is None:
            raise ValueError("请指定model和对应的keys，keys需要包含在model的字段内")
        self.index_name = self.index_name or self.index_name
        self.index_keys = self.index_keys or [k for k in self.keys if k != "id"]
        self.analyzer = RegexAnalyzer(r"([\u4e00-\u9fa5])|(\w+(\.?\w+)*)|[\uff08-\uff09]")
        self.schema = self.__get_schema()
        self.query = query
        self.queryset = queryset
        self.is_need_all = query is None and queryset is None
        # exists_id: {"创建索引的库": [id1, id2, id3]}
        self.path0exists_id = str(settings.RESOURCE_DIR / "/search/exists_id.json")
        self.path0whoosh_index = str(settings.RESOURCE_DIR / f"/search/{self.index_name}")

    def __get_schema(self):
        schema_para = {"id": ID(stored=True)}
        for field in self.model._meta.fields:
            if field.name != "id" and field.name in self.keys:
                if field.name in self.index_keys:
                    schema_para[field.name] = TEXT(stored=True, analyzer=self.analyzer)
                else:
                    schema_para[field.name] = KEYWORD(stored=True)
        return Schema(**schema_para)

    def __get_data(self) -> set:
        """
        现有的数据
        """
        need_update_pks = set()
        if self.queryset is None:
            self.queryset = self.model.objects.all().values(self.keys)
            self.queryset = self.queryset.filter(self.query)
        # 更新对应的id信息
        this_data = [k["id"] for k in self.queryset]
        with open(self.path0exists_id, "w") as fp:
            try:
                exists_data = json.load(fp)
            except Exception as e:
                print(e)
                exists_data = dict()
            if not self.is_need_all and exists_data:
                exists_pks = {pk for pk in exists_data[self.index_name]}
                need_update_pks = set(k for k in this_data if k in exists_pks)
            exists_data[self.index_name] = this_data
            fp.write(json.dumps(exists_data, ensure_ascii=False))
        return need_update_pks

    def create_index(self):
        """
        创建索引
        """
        if not os.path.exists(self.path0whoosh_index):
            os.mkdir(self.path0whoosh_index)
        need_update = self.__get_data() or set()
        storage = FileStorage(self.path0whoosh_index)
        try:
            ix = storage.open_index()
        except EmptyIndexError:
            ix = storage.create_index(self.schema)
        if self.is_need_all:
            ix.create_in()

        if not self.queryset:
            return
        writer = ix.writer()
        for data in self.queryset:
            if data["id"] in need_update:
                writer.update_document(**data)
            else:
                writer.add_document(**data)
        writer.commit()

    def filter(self, ix, **kwargs):
        """
        可以添加搜索条件
        """
        raise

    def search(self, page: int = 1, page_count: int = 20, **kwargs):
        """
        搜索
        :param page: 结果的第几页
        :param page_count: 每页有几个结果
        :return:
        """

        storage = FileStorage(self.path0whoosh_index)
        try:
            ix = storage.open_index()
        except EmptyIndexError:
            return [], 0
        query_search = self.filter(ix, **kwargs)
        if query_search is None:
            return {}, 0
        res = {}
        with ix.searcher() as s:
            result = s.search_page(query_search, pagelen=page_count, pagenum=page)
            total = result.total
            result.formatter = HtmlFormatter(tagname='em', classname='m', termclass='t')
            result.fragmenter = WholeFragmenter()
            for obj in result:
                key = obj.get('id')
                res[key] = obj.highlights('title')

        return res, total
